import numpy as np
import pandas as pd
from netZooPy.ligress.bonobo import Bonobo
from math import floor


def simulate_bonobo_mixture_varyDelta(sim_mean1, sim_cov1, sim_mean2, sim_cov2, mix_prop, delta, nsample = 100, nsample_ind = 100):
  sim_mean1 = sim_mean1["x"].to_numpy()
  sim_mean2 = sim_mean2["x"].to_numpy()
  
  size1 = floor(mix_prop*nsample)
  size2 = floor((1-mix_prop)*nsample)

  # Simulate individual gene expression and save individual mean and covariance matrix
  ind_mean = []
  ind_cor = []
  # Simulate from population1
  for i in range(size1):
    expr_sim = np.random.multivariate_normal(sim_mean1, sim_cov1, nsample_ind).T
    expr_mean = np.mean(expr_sim, 1)
    expr_cor = np.corrcoef(expr_sim)
    ind_mean.append(expr_mean)
    ind_cor.append(expr_cor)
  
  
  # Simulate from population2
  for i in range(size2):
    expr_sim = np.random.multivariate_normal(sim_mean2, sim_cov2, nsample_ind).T
    expr_mean = np.mean(expr_sim, 1)
    expr_cor = np.corrcoef(expr_sim)
    ind_mean.append(expr_mean)
    ind_cor.append(expr_cor)
  

  expression = pd.DataFrame(np.array(ind_mean).transpose())

  bonobo_obj = Bonobo(expression)
  bonobo_obj.run_bonobo(keep_in_memory=True, output_fmt='.txt', delta = delta)
  ind_bonobo = bonobo_obj.bonobos
  
  bonobo_obj = Bonobo(expression)
  bonobo_obj.run_bonobo(keep_in_memory=True, output_fmt='.txt', sparsify=True, delta = delta)
  ind_bonobo_sparse = bonobo_obj.bonobos


  # Compute RMSE for each bonobo
  mse_bonobo = []
  mse_bonobo_sparse = []
  for i in range(nsample):
    cov_error = (ind_cor[i] - ind_bonobo[i])**2
    cov_mse = np.mean(np.array(cov_error))
    mse_bonobo.append(cov_mse)
    
    cov_error = (ind_cor[i] - ind_bonobo_sparse[i])**2
    cov_mse = np.mean(np.array(cov_error))
    mse_bonobo_sparse.append(cov_mse)


  return mse_bonobo, mse_bonobo_sparse

